import pymysql
from typing import List, Dict, Any
from python_depend.mysql_pool import MySQLPool


# ==============================
# 数据库操作函数
# ==============================
def fetch_all_stock_list() -> List[Dict[str, Any]]:
    MySQLPool.initialize()
    mysql_pool = MySQLPool.get_pool('ainvestdb')
    """从数据库获取股票列表"""
    sql_query = """
    SELECT
        CAST(thsmarket_code_hq + 256 AS SIGNED) AS market,
        thscode_hq AS code
    FROM
        pub205
    WHERE isvalid = 1
        AND F005V_PUB205 IN ('212049', '212050', '212010', '212011', '212210')
        AND f018v_pub205 = '正常上市'
        AND F007D_PUB205 IS NOT NULL
        AND thscode_hq IS NOT NULL
        AND F003V_PUB205 IN ('海外普通股', '优先股', '存托凭证(普通股)', '存托凭证(优先股)', 'UNIT')
        order by  f019v_pub205 desc;
    """
    try:
        with mysql_pool.connection() as connection:
            with connection.cursor() as cursor:
                cursor.execute(sql_query)
                return cursor.fetchall()
    except pymysql.Error as e:
        print(f"[ERROR] 获取股票列表失败: {e}")
        return []
    MySQLPool.close_all()

def fetch_spx_stock_list() -> List[Dict[str, Any]]:
    MySQLPool.initialize()
    mysql_pool = MySQLPool.get_pool('ainvestdb')
    """从数据库获取股票列表"""
    sql_query = """
    SELECT
        CAST(thsmarket_code_hq + 256 AS SIGNED) AS market,
        thscode_hq AS code
    FROM
        pub205
    where isvalid = 1
    AND f018v_pub205 = '正常上市'
    AND F007D_PUB205 IS NOT NULL
    and PUB205.F019V_PUB205 in(
    select F001V_IND004 from  IND004 
    where ind004.isvalid = 1  and  F002V_IND004 = 'ind1028'
    and (IND004.F003 is null   or IND004.F003 >=current_DATE )
    ) order by  f019v_pub205 desc;
    """
    try:
        with mysql_pool.connection() as connection:
            with connection.cursor() as cursor:
                cursor.execute(sql_query)
                return cursor.fetchall()
    except pymysql.Error as e:
        print(f"[ERROR] 获取股票列表失败: {e}")
        return []
    MySQLPool.close_all()